from itertools import combinations

import numpy as np
import pandas as pd
import scipy.stats as stats
# https://mp.weixin.qq.com/s?__biz=MjM5ODc2Mzk2MA==&mid=2451789461&idx=4&sn=6f3f413bb140e831a5cd44aea225aaa3&chksm=b1123a428665b354fee71d06633c2381dcb479d1c63a13fbc8b9576665c58ae6fb0923c0c7f7&scene=27
# 斯皮尔曼相关系数（Spearman）也被叫做斯皮尔曼等级相关系数
# 计算两列的rank
# di = 两列的rank之差
# (di)^2
# n 为数据量，行数
# Rs = 1 - [6*E(di)^2)] / (n*(n^2 - 1))
# 标准化和rank相关预处理，并不会影响最终的Spearman结果
from sklearn.preprocessing import StandardScaler

pd_data = pd.DataFrame(
    [[1, 33, 40, 9876], [3434, 4, 4, 6], [12, 7, 222, 99], [2333, 34, 54, 22],
     [23, 56, 2, 987], [9.8, 3.4, 999.3, 0.9]],
    columns=['x_0', 'x_1', 'x_2', 'x_3'])

#
rank_value = stats.rankdata(pd_data, axis=0)
print(rank_value)

scaler = StandardScaler()
data_scale = scaler.fit_transform(pd_data)

pd_data_scale = pd.DataFrame(data_scale, columns=pd_data.columns)
print(pd_data_scale, f'\n{"*"*36}')
print('pd_data_scale>spearman:\n', pd_data_scale.corr('spearman'),
      f'\n{"*"*36}')


def spearman(data_frame: pd.DataFrame):
    print('data_raw>pearson:\n', data_frame.corr('pearson'), f'\n{"*"*36}')
    print('data_raw>spearman:\n', data_frame.corr('spearman'), f'\n{"*"*36}')
    data_rank = data_frame.rank()
    print(data_rank)

    print('data_rank>pearson:\n', data_rank.corr('pearson'), f'\n{"*"*36}')
    print('data_rank>spearman:\n', data_rank.corr('spearman'), f'\n{"*"*36}')

    data_headers = data_rank.columns.values

    data_headers_mapper = {i: idx for idx, i in enumerate(data_headers)}

    N, K = data_rank.shape
    correl = np.empty((K, K), dtype=float)

    for i in range(len(correl)):
        correl[i][i] = 1

    for h1, h2 in combinations(data_headers, 2):
        d = data_rank[h1] - data_rank[h2]
        d2 = d**2
        d2_sum = d2.sum()
        rs_val = (1 - (6 * d2_sum) / (N * (N**2 - 1)))
        correl[data_headers_mapper[h1]][data_headers_mapper[h2]] = rs_val
        correl[data_headers_mapper[h2]][data_headers_mapper[h1]] = rs_val

    return pd.DataFrame(correl, index=data_headers, columns=data_headers)


print(spearman(pd_data))
